package common

import (
	"bytes"
	"fmt"
	"golang.org/x/crypto/ssh"
	"io/ioutil"
	"minissh/flags"
	"net"
	"strconv"
	"time"
)

//func Connect(user, password, host, key string, port int, cipherList, keyExchangeList []string) (*ssh.Session, error) {
func Connect(c *Conf) (*ssh.Session, error) {

	var (
		auth         []ssh.AuthMethod
		addr         string
		config       ssh.Config
		clientConfig *ssh.ClientConfig
		client       *ssh.Client
		session      *ssh.Session
		err          error
	)
	auth = make([]ssh.AuthMethod, 0)
	if c.Key == "" {
		auth = append(auth, ssh.Password(c.Password))
	} else {
		pemBytes, err := ioutil.ReadFile(c.Key)
		if err != nil {
			return nil, err
		}

		var signer ssh.Signer
		if c.Password == "" {
			signer, err = ssh.ParsePrivateKey(pemBytes)
		} else {
			signer, err = ssh.ParsePrivateKeyWithPassphrase(pemBytes, []byte(c.Password))
		}
		if err != nil {
			return nil, err
		}
		auth = append(auth, ssh.PublicKeys(signer))
	}

	if len(c.CipherList) == 0 {
		config.Ciphers = []string{"aes128-ctr", "aes192-ctr", "aes256-ctr", "aes128-gcm@openssh.com", "arcfour256", "arcfour128", "aes128-cbc", "3des-cbc", "aes192-cbc", "aes256-cbc"}
	} else {
		config.Ciphers = c.CipherList
	}

	if len(c.KeyExchangeList) == 0 {
		config.KeyExchanges = []string{"diffie-hellman-group-exchange-sha1", "diffie-hellman-group1-sha1", "diffie-hellman-group-exchange-sha256"}
	} else {
		config.KeyExchanges = c.KeyExchangeList
	}

	clientConfig = &ssh.ClientConfig{
		Auth:    auth,
		User:    c.User,
		Timeout: time.Second * 30,
		Config:  config,
		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
			return nil
		},
	}

	//连接ssh
	addr = fmt.Sprintf("%s:%d", c.Host, c.Port)
	if client, err = ssh.Dial("tcp", addr, clientConfig); err != nil {
		return nil, err
	}

	//create session
	if session, err = client.NewSession(); err != nil {
		return nil, err
	}

	models := ssh.TerminalModes{
		ssh.ECHO:          0,     // disable echoing
		ssh.TTY_OP_ISPEED: 14400, // input speed = 14.4kbaud
		ssh.TTY_OP_OSPEED: 14400, // output speed = 14.4kbaud
	}
	if err = session.RequestPty("xterm", 80, 40, models); err != nil {
		return nil, err
	}

	return session, nil
}

func DoSSH(c *Conf, cf *flags.Flags, ch chan SSHResult) {
	chSSH := make(chan SSHResult)
	if cf.LinuxMode {
		go doSSH_Run(c, cf.Cmds, ch)
	} else {
		go doSSH_Session(c, cf.Cmds, ch)
	}

	var res SSHResult

	select {
	case <-time.After(time.Duration(cf.Timeout) * time.Second):
		res.Host = c.Host
		res.Success = false
		res.Result = ("SSH run timeout：" + strconv.Itoa(cf.Timeout) + " second.")
		ch <- res
	case res = <-chSSH:
		ch <- res
	}
	return
}

func doSSH_Session(c *Conf, cmd string, ch chan SSHResult) {
	session, err := Connect(c)
	var sshResult SSHResult
	sshResult.Host = c.Host

	if err != nil {
		sshResult.Success = false
		sshResult.Result = fmt.Sprintf("<%s>", err.Error())
		ch <- sshResult
		return
	}
	defer session.Close()

	stdinBuf, _ := session.StdinPipe()

	var outBytes, errBytes bytes.Buffer
	session.Stdout = &outBytes
	session.Stderr = &errBytes

	err = session.Shell()
	if err != nil {
		sshResult.Success = false
		sshResult.Result = fmt.Sprintf("<%s>", err.Error())
		ch <- sshResult
		return
	}
	stdinBuf.Write([]byte(cmd))

	session.Wait()
	if errBytes.String() != "" {
		sshResult.Success = false
		sshResult.Result = errBytes.String()
		ch <- sshResult
	} else {
		sshResult.Success = true
		sshResult.Result = outBytes.String()
		ch <- sshResult
	}
	return
}

func doSSH_Run(c *Conf, cmd string, ch chan SSHResult) {
	session, err := Connect(c)
	var sshResult SSHResult
	sshResult.Host = c.Host

	if err != nil {
		sshResult.Success = false
		sshResult.Result = fmt.Sprintf("<%s>", err.Error())
		ch <- sshResult
		return
	}
	defer session.Close()

	var outBytes, errBytes bytes.Buffer
	session.Stdout = &outBytes
	session.Stderr = &errBytes

	err = session.Run(cmd)
	if err != nil {
		sshResult.Success = false
		sshResult.Result = fmt.Sprintf("<%s>", err.Error())
		ch <- sshResult
		return
	}

	if errBytes.String() != "" {
		sshResult.Success = false
		sshResult.Result = errBytes.String()
		ch <- sshResult
	} else {
		sshResult.Success = true
		sshResult.Result = outBytes.String()
		ch <- sshResult
	}
	return
}
